# 2488. 统计中位数为 K 的子数组
# 给你一个长度为 n 的数组 nums ，该数组由从 1 到 n 的 不同 整数组成。另给你一个正整数 k 。

# 统计并返回 nums 中的 中位数 等于 k 的非空子数组的数目。

# 注意：

# 数组的中位数是按 递增 顺序排列后位于 中间 的那个元素，如果数组长度为偶数，则中位数是位于中间靠 左 的那个元素。
# 例如，[2,3,1,4] 的中位数是 2 ，[8,4,3,5,1] 的中位数是 4 。
# 子数组是数组中的一个连续部分。
#  

# 示例 1：

# 输入：nums = [3,2,1,4,5], k = 4
# 输出：3
# 解释：中位数等于 4 的子数组有：[4]、[4,5] 和 [1,4,5] 。
# 示例 2：

# 输入：nums = [2,3,1], k = 3
# 输出：1
# 解释：[3] 是唯一一个中位数等于 3 的子数组。
#  

# 提示：

# n == nums.length
# 1 <= n <= 105
# 1 <= nums[i], k <= n
# nums 中的整数互不相同

# 来源：力扣（LeetCode）
# 链接：https://leetcode.cn/problems/count-subarrays-with-median-k
# 著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。


from time import time
from typing import List
from collections import Counter


class Solution:
    
    def sign(self, num: int) -> int:
        return (num > 0) - (num < 0)

    def countSubarrays(self, nums: List[int], k: int) -> int:
        kIndex = nums.index(k)
        ans = 0
        counts = Counter()
        counts[0] = 1
        sum = 0
        for i, num in enumerate(nums):
            sum += self.sign(num - k)
            if i < kIndex:
                counts[sum] += 1
            else:
                prev0 = counts[sum]
                prev1 = counts[sum - 1]
                ans += prev0 + prev1
        return ans
    
    
if __name__ == '__main__':
    # args = { "nums": [3,2,1,4,5], "k": 4}
    args = { "nums": [2,3,1], "k": 3}
    start = time()
    print(Solution().countSubarrays(**args))
    print('='*40)
    print('耗时:', time()*1000 - start*1000, 'ms')